# -*- coding: utf-8 -*-
"""
数据库操作模块
"""
import logging
import pymysql
import pandas as pd
from sqlalchemy import create_engine, text
from typing import List, Dict, Any, Optional, Tuple
from contextlib import contextmanager
from datetime import datetime
from config import DATABASE_CONFIG

logger = logging.getLogger(__name__)

class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, config: Dict[str, Any] = None):
        self.config = config or DATABASE_CONFIG
        # 创建SQLAlchemy引擎
        self.engine = self._create_engine()
    
    def _create_engine(self):
        """创建SQLAlchemy引擎"""
        connection_string = (
            f"mysql+pymysql://{self.config['user']}:{self.config['password']}"
            f"@{self.config['host']}:{self.config['port']}/{self.config['database']}"
            f"?charset={self.config['charset']}"
        )
        return create_engine(connection_string, echo=False)
    
    @contextmanager
    def get_connection(self):
        """获取数据库连接"""
        connection = None
        try:
            connection = pymysql.connect(**self.config)
            yield connection
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            if connection:
                connection.rollback()
            raise
        finally:
            if connection:
                connection.close()
    
    def execute_query(self, sql: str, params: Tuple = None) -> List[Dict[str, Any]]:
        """执行查询语句"""
        with self.get_connection() as conn:
            cursor = conn.cursor(pymysql.cursors.DictCursor)
            try:
                cursor.execute(sql, params)
                return cursor.fetchall()
            finally:
                cursor.close()
    
    def execute_update(self, sql: str, params: Tuple = None) -> int:
        """执行更新语句"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            try:
                affected_rows = cursor.execute(sql, params)
                conn.commit()
                return affected_rows
            except Exception as e:
                conn.rollback()
                logger.error(f"执行更新失败: {e}")
                raise
            finally:
                cursor.close()
    
    def execute_batch_update(self, sql: str, params_list: List[Tuple]) -> int:
        """批量执行更新语句"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            try:
                affected_rows = cursor.executemany(sql, params_list)
                conn.commit()
                return affected_rows
            except Exception as e:
                conn.rollback()
                logger.error(f"批量更新失败: {e}")
                raise
            finally:
                cursor.close()

class StockInfoDAO:
    """股票信息数据访问对象"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db = db_manager
    
    def insert_or_update_stock_info(self, stock_data: Dict[str, Any]) -> int:
        """插入或更新股票信息"""
        sql = """
        INSERT INTO new_stock_info 
        (stock_code, stock_name, exchange, industry, market, listing_date, status)
        VALUES (%s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        stock_name = VALUES(stock_name),
        exchange = VALUES(exchange),
        industry = VALUES(industry),
        market = VALUES(market),
        listing_date = VALUES(listing_date),
        status = VALUES(status),
        updated_at = CURRENT_TIMESTAMP
        """
        params = (
            stock_data['stock_code'],
            stock_data['stock_name'],
            stock_data['exchange'],
            stock_data.get('industry'),
            stock_data.get('market'),
            stock_data.get('listing_date'),
            stock_data.get('status', 'active')
        )
        return self.db.execute_update(sql, params)
    
    def batch_insert_or_update_stock_info(self, stock_data_list: List[Dict[str, Any]]) -> int:
        """批量插入或更新股票信息"""
        sql = """
        INSERT INTO new_stock_info 
        (stock_code, stock_name, exchange, industry, market, listing_date, status)
        VALUES (%s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        stock_name = VALUES(stock_name),
        exchange = VALUES(exchange),
        industry = VALUES(industry),
        market = VALUES(market),
        listing_date = VALUES(listing_date),
        status = VALUES(status),
        updated_at = CURRENT_TIMESTAMP
        """
        params_list = []
        for stock_data in stock_data_list:
            params_list.append((
                stock_data['stock_code'],
                stock_data['stock_name'],
                stock_data['exchange'],
                stock_data.get('industry'),
                stock_data.get('market'),
                stock_data.get('listing_date'),
                stock_data.get('status', 'active')
            ))
        return self.db.execute_batch_update(sql, params_list)
    
    def get_all_stock_codes(self) -> List[str]:
        """获取所有股票代码"""
        sql = "SELECT stock_code FROM new_stock_info WHERE status = 'active'"
        results = self.db.execute_query(sql)
        return [row['stock_code'] for row in results]
    
    def get_stock_info_by_code(self, stock_code: str) -> Optional[Dict[str, Any]]:
        """根据股票代码获取股票信息"""
        sql = "SELECT * FROM new_stock_info WHERE stock_code = %s"
        results = self.db.execute_query(sql, (stock_code,))
        return results[0] if results else None
    
    def batch_insert_or_update_with_pandas(self, df: pd.DataFrame) -> int:
        """使用pandas批量插入或更新股票信息"""
        try:
            # 使用pandas的to_sql方法，配合SQLAlchemy
            affected_rows = df.to_sql(
                'new_stock_info',
                self.db.engine,
                if_exists='append',
                index=False,
                chunksize=1000,
                method='multi'
            )
            logger.info(f"使用pandas批量插入 {len(df)} 条股票信息")
            return len(df)
        except Exception as e:
            logger.error(f"pandas批量插入失败: {e}")
            # 如果pandas插入失败，回退到传统方法
            return self.batch_insert_or_update_stock_info(df.to_dict('records'))

class TagDAO:
    """标签数据访问对象"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db = db_manager
    
    def insert_or_update_tag(self, tag_data: Dict[str, Any]) -> int:
        """插入或更新标签"""
        sql = """
        INSERT INTO tags (name, type, source)
        VALUES (%s, %s, %s)
        ON DUPLICATE KEY UPDATE
        source = VALUES(source),
        updated_at = CURRENT_TIMESTAMP
        """
        params = (
            tag_data['name'],
            tag_data['type'],
            tag_data.get('source', 'akshare')
        )
        return self.db.execute_update(sql, params)
    
    def get_tag_id(self, name: str, type: str) -> Optional[int]:
        """获取标签ID"""
        sql = "SELECT tag_id FROM tags WHERE name = %s AND type = %s"
        results = self.db.execute_query(sql, (name, type))
        return results[0]['tag_id'] if results else None
    
    def insert_or_update_stock_tag(self, stock_code: str, tag_id: int, source: str = 'akshare') -> int:
        """插入或更新股票标签关联"""
        sql = """
        INSERT INTO stock_tag (code, tag_id, source)
        VALUES (%s, %s, %s)
        ON DUPLICATE KEY UPDATE
        source = VALUES(source),
        updated_at = CURRENT_TIMESTAMP
        """
        params = (stock_code, tag_id, source)
        return self.db.execute_update(sql, params)
    
    def batch_insert_or_update_stock_tags(self, stock_tag_list: List[Dict[str, Any]]) -> int:
        """批量插入或更新股票标签关联"""
        sql = """
        INSERT INTO stock_tag (code, tag_id, source)
        VALUES (%s, %s, %s)
        ON DUPLICATE KEY UPDATE
        source = VALUES(source),
        updated_at = CURRENT_TIMESTAMP
        """
        params_list = []
        for item in stock_tag_list:
            params_list.append((
                item['code'],
                item['tag_id'],
                item.get('source', 'akshare')
            ))
        return self.db.execute_batch_update(sql, params_list)
    
    def get_stock_tags(self, stock_code: str) -> List[Dict[str, Any]]:
        """获取股票的标签信息"""
        sql = """
        SELECT t.tag_id, t.name, t.type, t.source, st.updated_at
        FROM stock_tag st
        INNER JOIN tags t ON st.tag_id = t.tag_id
        WHERE st.code = %s
        ORDER BY t.type, t.name
        """
        return self.db.execute_query(sql, (stock_code,))
    
    def get_tags_by_type(self, tag_type: str) -> List[Dict[str, Any]]:
        """根据类型获取标签列表"""
        sql = "SELECT * FROM tags WHERE type = %s ORDER BY name"
        return self.db.execute_query(sql, (tag_type,))
    
    def delete_stock_tags(self, stock_code: str, tag_type: str = None) -> int:
        """删除股票的标签关联"""
        if tag_type:
            sql = """
            DELETE st FROM stock_tag st
            INNER JOIN tags t ON st.tag_id = t.tag_id
            WHERE st.code = %s AND t.type = %s
            """
            return self.db.execute_update(sql, (stock_code, tag_type))
        else:
            sql = "DELETE FROM stock_tag WHERE code = %s"
            return self.db.execute_update(sql, (stock_code,))

class StockDailyDAO:
    """股票日线数据访问对象"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db = db_manager
    
    def insert_or_update_daily_data(self, daily_data: Dict[str, Any]) -> int:
        """插入或更新日线数据"""
        sql = """
        INSERT INTO new_stock_daily 
        (stock_code, trade_date, open_price, high_price, low_price, close_price, volume, amount, 
         amplitude, change_pct, change_amount, turnover_rate)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        open_price = VALUES(open_price),
        high_price = VALUES(high_price),
        low_price = VALUES(low_price),
        close_price = VALUES(close_price),
        volume = VALUES(volume),
        amount = VALUES(amount),
        amplitude = VALUES(amplitude),
        change_pct = VALUES(change_pct),
        change_amount = VALUES(change_amount),
        turnover_rate = VALUES(turnover_rate),
        updated_at = CURRENT_TIMESTAMP
        """
        params = (
            daily_data['stock_code'],
            daily_data['trade_date'],
            daily_data['open_price'],
            daily_data['high_price'],
            daily_data['low_price'],
            daily_data['close_price'],
            daily_data['volume'],
            daily_data.get('amount'),
            daily_data.get('amplitude'),
            daily_data.get('change_pct'),
            daily_data.get('change_amount'),
            daily_data.get('turnover_rate')
        )
        return self.db.execute_update(sql, params)
    
    def batch_insert_or_update_daily_data(self, daily_data_list: List[Dict[str, Any]]) -> int:
        """批量插入或更新日线数据"""
        sql = """
        INSERT INTO new_stock_daily 
        (stock_code, trade_date, open_price, high_price, low_price, close_price, volume, amount, 
         amplitude, change_pct, change_amount, turnover_rate)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        open_price = VALUES(open_price),
        high_price = VALUES(high_price),
        low_price = VALUES(low_price),
        close_price = VALUES(close_price),
        volume = VALUES(volume),
        amount = VALUES(amount),
        amplitude = VALUES(amplitude),
        change_pct = VALUES(change_pct),
        change_amount = VALUES(change_amount),
        turnover_rate = VALUES(turnover_rate),
        updated_at = CURRENT_TIMESTAMP
        """
        params_list = []
        for daily_data in daily_data_list:
            params_list.append((
                daily_data['stock_code'],
                daily_data['trade_date'],
                daily_data['open_price'],
                daily_data['high_price'],
                daily_data['low_price'],
                daily_data['close_price'],
                daily_data['volume'],
                daily_data.get('amount'),
                daily_data.get('amplitude'),
                daily_data.get('change_pct'),
                daily_data.get('change_amount'),
                daily_data.get('turnover_rate')
            ))
        return self.db.execute_batch_update(sql, params_list)
    
    def get_latest_trade_date(self, stock_code: str) -> Optional[str]:
        """获取指定股票的最新交易日期"""
        sql = """
        SELECT trade_date FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT 1
        """
        results = self.db.execute_query(sql, (stock_code,))
        return results[0]['trade_date'].strftime('%Y-%m-%d') if results else None
    
    def get_daily_data_by_date_range(self, stock_code: str, start_date: str, end_date: str) -> List[Dict[str, Any]]:
        """获取指定日期范围的日线数据"""
        sql = """
        SELECT * FROM new_stock_daily 
        WHERE stock_code = %s AND trade_date BETWEEN %s AND %s
        ORDER BY trade_date
        """
        return self.db.execute_query(sql, (stock_code, start_date, end_date))

class StockMinuteDAO:
    """股票分钟级数据访问对象"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db = db_manager
    
    def insert_or_update_minute_data(self, minute_data: Dict[str, Any]) -> int:
        """插入或更新分钟级数据"""
        sql = """
        INSERT INTO new_stock_minute_data 
        (stock_code, market, trade_time, open, close, high, low, volume, amount)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        open = VALUES(open),
        close = VALUES(close),
        high = VALUES(high),
        low = VALUES(low),
        volume = VALUES(volume),
        amount = VALUES(amount),
        update_time = CURRENT_TIMESTAMP
        """
        params = (
            minute_data['stock_code'],
            minute_data['market'],
            minute_data['trade_time'],
            minute_data['open'],
            minute_data['close'],
            minute_data['high'],
            minute_data['low'],
            minute_data['volume'],
            minute_data['amount']
        )
        return self.db.execute_update(sql, params)
    
    def batch_insert_or_update_minute_data(self, minute_data_list: List[Dict[str, Any]]) -> int:
        """批量插入或更新分钟级数据"""
        sql = """
        INSERT INTO new_stock_minute_data 
        (stock_code, market, trade_time, open, close, high, low, volume, amount)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
        open = VALUES(open),
        close = VALUES(close),
        high = VALUES(high),
        low = VALUES(low),
        volume = VALUES(volume),
        amount = VALUES(amount),
        update_time = CURRENT_TIMESTAMP
        """
        params_list = []
        for minute_data in minute_data_list:
            params_list.append((
                minute_data['stock_code'],
                minute_data['market'],
                minute_data['trade_time'],
                minute_data['open'],
                minute_data['close'],
                minute_data['high'],
                minute_data['low'],
                minute_data['volume'],
                minute_data['amount']
            ))
        return self.db.execute_batch_update(sql, params_list)
    
    def get_latest_trade_time(self, stock_code: str, market: str) -> Optional[datetime]:
        """获取指定股票和市场的最新交易时间"""
        sql = """
        SELECT trade_time FROM new_stock_minute_data 
        WHERE stock_code = %s AND market = %s 
        ORDER BY trade_time DESC 
        LIMIT 1
        """
        results = self.db.execute_query(sql, (stock_code, market))
        return results[0]['trade_time'] if results else None
    
    def get_minute_data_by_time_range(self, stock_code: str, market: str, 
                                    start_time: datetime, end_time: datetime) -> List[Dict[str, Any]]:
        """获取指定时间范围的分钟级数据"""
        sql = """
        SELECT * FROM new_stock_minute_data 
        WHERE stock_code = %s AND market = %s AND trade_time BETWEEN %s AND %s
        ORDER BY trade_time
        """
        return self.db.execute_query(sql, (stock_code, market, start_time, end_time))
    
    def batch_insert_or_update_with_pandas(self, df: pd.DataFrame) -> int:
        """使用pandas批量插入或更新分钟级数据"""
        try:
            # 使用pandas的to_sql方法，配合SQLAlchemy
            affected_rows = df.to_sql(
                'new_stock_minute_data',
                self.db.engine,
                if_exists='append',
                index=False,
                chunksize=1000,
                method='multi'
            )
            logger.info(f"使用pandas批量插入 {len(df)} 条分钟级数据")
            return len(df)
        except Exception as e:
            logger.error(f"pandas批量插入分钟级数据失败: {e}")
            # 如果pandas插入失败，回退到传统方法
            return self.batch_insert_or_update_minute_data(df.to_dict('records'))
    
    def get_data_statistics(self) -> Dict[str, Any]:
        """获取分钟级数据统计信息"""
        try:
            # 获取总记录数
            total_records = self.db.execute_query('SELECT COUNT(*) as count FROM new_stock_minute_data')[0]['count']
            
            # 获取股票数量
            total_stocks = self.db.execute_query('SELECT COUNT(DISTINCT stock_code) as count FROM new_stock_minute_data')[0]['count']
            
            # 获取最新和最旧的时间
            latest_time = self.db.execute_query('SELECT MAX(trade_time) as time FROM new_stock_minute_data')[0]['time']
            earliest_time = self.db.execute_query('SELECT MIN(trade_time) as time FROM new_stock_minute_data')[0]['time']
            
            return {
                'total_records': total_records,
                'total_stocks': total_stocks,
                'latest_time': latest_time.strftime('%Y-%m-%d %H:%M:%S') if latest_time else None,
                'earliest_time': earliest_time.strftime('%Y-%m-%d %H:%M:%S') if earliest_time else None
            }
        except Exception as e:
            logger.error(f"获取分钟级数据统计失败: {e}")
            return {
                'error': str(e),
                'total_records': 0,
                'total_stocks': 0,
                'latest_time': None,
                'earliest_time': None
            }
